# coding='utf-8'
"""

（多）分类问题下的指标评估（acc, precision, recall, f1）。

"""
from typing import List

import numpy as np  # 导入 NumPy，用于矩阵运算
import pandas as pd  # 导入 Pandas，虽然未使用，但可能是预留功能
from sklearn.metrics import accuracy_score, precision_score, f1_score, recall_score, confusion_matrix  # 导入评估指标


class ClassEvaluator(object):
    def __init__(self):
        """
        初始化评估器，累积预测和真实标签。
        """
        self.goldens = []  # 存储真实标签列表
        self.predictions = []  # 存储预测标签列表

    def add_batch(self, pred_batch: List[List], gold_batch: List[List]):
        """
        添加一个batch中的prediction和gold列表，用于后续统一计算。

        Args:
            pred_batch (list): 模型预测标签列表, e.g. ->  [['体', '育'], ['财', '经'], ...]
            gold_batch (list): 真实标签标签列表, e.g. ->  [['体', '育'], ['财', '经'], ...]
        """
        assert len(pred_batch) == len(gold_batch), "预测标签和真实标签长度不一致！"
        if type(gold_batch[0]) in [list, tuple]:  # 若遇到多个子标签构成一个标签的情况
            pred_batch = [''.join([str(e) for e in ele]) for ele in pred_batch]  # 将所有子标签拼接为完整标签
            gold_batch = [''.join([str(e) for e in ele]) for ele in gold_batch]
        self.goldens.extend(gold_batch)  # 将当前 batch 的真实标签添加到累积列表中
        self.predictions.extend(pred_batch)  # 将当前 batch 的预测标签添加到累积列表中

    def compute(self, round_num=2) -> dict:
        """
        根据当前类中累积的变量值，计算当前的P, R, F1。

        Args:
            round_num (int): 计算结果保留小数点后几位, 默认小数点后2位。

        Returns:
            dict -> {
                'accuracy': 准确率,
                'precision': 精准率,
                'recall': 召回率,
                'f1': f1值,
                'class_metrics': {
                    '0': {
                            'precision': 该类别下的precision,
                            'recall': 该类别下的recall,
                            'f1': 该类别下的f1
                        },
                    ...
                }
            }
        """
        classes = sorted(list(set(self.goldens) | set(self.predictions)))  # 获取所有唯一类别
        class_metrics, res = {}, {}

        # 构建全局指标：准确率、精准率、召回率、F1值
        res['accuracy'] = round(accuracy_score(self.goldens, self.predictions), round_num)
        res['precision'] = round(precision_score(self.goldens, self.predictions, average='weighted'), round_num)
        res['recall'] = round(recall_score(self.goldens, self.predictions, average='weighted'), round_num)
        res['f1'] = round(f1_score(self.goldens, self.predictions, average='weighted'), round_num)

        try:
            conf_matrix = np.array(confusion_matrix(self.goldens, self.predictions))  # 计算混淆矩阵 (n_class, n_class)
            assert conf_matrix.shape[0] == len(classes), "混淆矩阵维度与类别数不一致！"
            for i in range(conf_matrix.shape[0]):  # 遍历每个类别
                precision = 0 if sum(conf_matrix[:, i]) == 0 else conf_matrix[i, i] / sum(conf_matrix[:, i])
                recall = 0 if sum(conf_matrix[i, :]) == 0 else conf_matrix[i, i] / sum(conf_matrix[i, :])
                f1 = 0 if (precision + recall) == 0 else 2 * precision * recall / (precision + recall)
                class_metrics[classes[i]] = {  # 构建每个类别的指标
                    'precision': round(precision, round_num),
                    'recall': round(recall, round_num),
                    'f1': round(f1, round_num)
                }
            res['class_metrics'] = class_metrics  # 将类别指标添加到结果字典中
        except Exception as e:
            print(f'[Warning] Something wrong when calculate class_metrics: {e}')
            print(f'-> goldens: {set(self.goldens)}')
            print(f'-> predictions: {set(self.predictions)}')
            print(f'-> diff elements: {set(self.predictions) - set(self.goldens)}')
            res['class_metrics'] = {}  # 如果出现异常，返回空类别指标

        return res  # 返回最终的评估结果

    def reset(self):
        """
        重置积累的数值。
        """
        self.goldens = []  # 清空真实标签列表
        self.predictions = []  # 清空预测标签列表


if __name__ == '__main__':
    from rich import print  # 导入 rich 库，用于美化打印输出

    # 实例化评估器
    metric = ClassEvaluator()

    # 添加批量数据
    metric.add_batch(
        [['财', '经'], ['财', '经'], ['体', '育'], ['体', '育'], ['计', '算', '机']],  # 预测标签
        [['体', '育'], ['财', '经'], ['体', '育'], ['计', '算', '机'], ['计', '算', '机']]  # 真实标签
    )
    # metric.add_batch(
    #     [0, 0, 1, 1, 0],
    #     [1, 1, 1, 0, 0]
    # )
    # 打印评估结果
    print(metric.compute())
# 这段代码实现了一个灵活且通用的分类评估器，适用于多标签或多分类任务。核心功能包括：

# 数据累积：支持按批次添加预测和真实标签，并自动处理多标签场景。
# 全局指标计算：基于 sklearn.metrics 提供的标准方法，计算准确率、精准率、召回率和 F1 值。
# 类别指标计算：通过混淆矩阵，计算每个类别的精准率、召回率和 F1 值，提供更细粒度的评估。
# 异常处理：在类别指标计算中加入异常捕获机制，确保代码的鲁棒性。
# 整体设计简洁高效，适用于自然语言处理、图像分类等任务中的模型性能评估。